'''
Author: chyang0822 270917365@qq.com
Date: 2024-08-22 15:51:42
LastEditors: chyang0822 270917365@qq.com
LastEditTime: 2024-08-22 15:51:45
FilePath: /Project-DMAI-Deploy/tutorials/tensorrt/1.9-int8/int8.py
Description: 

Copyright (c) 2024 by ${git_name_email}, All Rights Reserved. 
'''
import os
import onnx
from tqdm import tqdm
import numpy as np
import seaborn as sns
from onnx import numpy_helper
import  matplotlib.pyplot as plt

quant=8
thresh = 2**quant//2

model_path = "build/classifier.onnx"
model_name = os.path.splitext(os.path.basename(model_path))[0]
model = onnx.load(model_path)
onnx_weights = {}
for initializer in model.graph.initializer:
    onnx_weights[initializer.name]=numpy_helper.to_array(initializer)
new_onnx_weight={}
weight_names=onnx_weights.keys()
for node in model.graph.node:
    name = node.name
    for ins in node.input:
        if ins in weight_names:
            new_name= name+"/"+ins
            new_onnx_weight[new_name]=onnx_weights[ins]
            
print(f"all the unstable range:")
unstable={}
for k,v in new_onnx_weight.items():
    max_ = np.round(np.max(v),2)
    min_ = np.round(np.min(v),2)
    txt = f"{k} num:{v.size} max:{max_} min:{min_}"
    if np.abs(min_)>thresh or np.abs(max_)>thresh or np.abs(max_-min_)>2*thresh:
        unstable[k]=txt
        print(txt)


num_columns=3
dpi=100
num_rows= len(new_onnx_weight)//num_columns + 1
fig,ax = plt.subplots(nrows=num_rows,ncols=num_columns,figsize=(8*num_columns,6*num_rows),dpi=dpi)

for i,(k,v) in tqdm(enumerate(new_onnx_weight.items()),total=len(new_onnx_weight)):
    row = i//num_columns
    col = i%num_columns
    sns.histplot(v.flatten(),kde=True,ax=ax[row][col],fill=True,color='yellow')
    max_ = np.round(np.max(v),2)
    min_ = np.round(np.min(v),2)
    txt = f"{k} num:{v.size} max:{str(max_)} min:{str(min_)}"
    ax[row][col].set_title(txt)
    #ax[row][col].set_facecolor('lightgray')
    if np.abs(min_)>thresh or np.abs(max_)>thresh or np.abs(max_-min_)>2*thresh:
        unstable[k]=txt
# fig.set_size_inches(6*num_columns/dpi,6*num_rows/dpi)
fig.savefig(f'{model_name}.jpg',bbox_inches='tight',pad_inches=0.1)
